#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建测试数据用于瞬变源分类程序测试
"""

import numpy as np
from astropy.io import fits
import os
import json


def create_test_fits(filename: str, image_size: int = 100):
    """
    创建测试FITS文件，包含观测图像、模板图像和残差图像
    
    参数:
        filename: 输出文件名
        image_size: 图像大小
    """
    # 创建基础背景
    base_background = 100
    noise_level = 10
    
    # 观测图像：包含一些模拟源
    obs_img = np.random.normal(base_background, noise_level, (image_size, image_size))
    
    # 模板图像：与观测图像相似但可能缺少某些源
    tpl_img = np.random.normal(base_background, noise_level, (image_size, image_size))
    
    # 残差图像：观测减去模板
    res_img = obs_img - tpl_img
    
    # 在观测图像中心添加一个模拟瞬变源
    center_x, center_y = image_size // 2, image_size // 2
    obs_img[center_y-2:center_y+3, center_x-2:center_x+3] += 50
    res_img[center_y-2:center_y+3, center_x-2:center_x+3] += 50
    
    # 在另一个位置添加热像素
    hot_x, hot_y = 30, 30
    obs_img[hot_y, hot_x] += 200  # 单个像素高亮
    res_img[hot_y, hot_x] += 200
    
    # 创建FITS文件
    hdul = fits.HDUList()
    
    # 主HDU（观测图像）
    primary_hdu = fits.PrimaryHDU(obs_img)
    primary_hdu.header['IMAGE_TYPE'] = 'OBJT'
    primary_hdu.header['COMMENT'] = 'Observation Image'
    hdul.append(primary_hdu)
    
    # 模板图像HDU
    tpl_hdu = fits.ImageHDU(tpl_img)
    tpl_hdu.header['IMAGE_TYPE'] = 'TEMP'
    tpl_hdu.header['COMMENT'] = 'Template Image'
    hdul.append(tpl_hdu)
    
    # 残差图像HDU
    res_hdu = fits.ImageHDU(res_img)
    res_hdu.header['IMAGE_TYPE'] = 'DIFF'
    res_hdu.header['COMMENT'] = 'Difference Image'
    hdul.append(res_hdu)
    
    # 保存文件
    hdul.writeto(filename, overwrite=True)
    print(f"创建测试FITS文件: {filename}")


def create_transient_list(output_file: str, num_objects: int = 5):
    """
    创建瞬变源列表JSON文件
    
    参数:
        output_file: 输出文件名
        num_objects: 对象数量
    """
    objects = []
    
    for i in range(num_objects):
        obj = {
            'name': f'transient_{i+1:03d}',
            'ra': 120.0 + i * 0.1,
            'dec': 30.0 + i * 0.1,
            'mag': 15.0 + i * 0.5
        }
        objects.append(obj)
    
    data = {
        'objects': objects,
        'obj_number': num_objects,
        'description': '测试瞬变源列表'
    }
    
    with open(output_file, 'w') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    
    print(f"创建瞬变源列表: {output_file}")


def setup_test_environment():
    """设置测试环境"""
    # 创建测试目录结构
    test_dir = 'data/datasets/test_20251004'
    images_dir = os.path.join(test_dir, 'images')
    
    os.makedirs(images_dir, exist_ok=True)
    
    # 创建瞬变源列表
    list_file = os.path.join(test_dir, 'transient_list.json')
    create_transient_list(list_file, 3)
    
    # 创建测试FITS文件
    with open(list_file, 'r') as f:
        data = json.load(f)
    
    for obj in data['objects']:
        name = obj['name']
        fits_file = os.path.join(images_dir, f"{name}_1.fit")
        create_test_fits(fits_file)
    
    print("测试环境设置完成")
    return list_file, images_dir


if __name__ == "__main__":
    # 创建单个测试FITS文件
    create_test_fits('test_image.fit')
    
    # 设置批量测试环境
    list_file, images_dir = setup_test_environment()
    
    print("\n测试数据创建完成，可以运行以下命令进行测试:")
    print("1. 单个图像分类:")
    print("   python transient_classifier.py")
    print("\n2. 批量分类:")
    print(f"   python -c \"from transient_classifier import batchClassify; batchClassify('{list_file}', '{images_dir}')\"")
    print("\n3. 或者直接运行:")
    print("   python")
    print("   >>> from transient_classifier import batchClassify")
    print(f"   >>> batchClassify('{list_file}', '{images_dir}')")